#include "kernel_operator.h"
#include <type_traits>
using namespace AscendC;
constexpr int32_t BUFFER_NUM = 2;
constexpr int32_t ACCTEST_NUM = 6;
class KernelScatterReduce {
public:
    __aicore__ inline KernelScatterReduce() {}
    __aicore__ inline void Init(GM_ADDR self, GM_ADDR index, GM_ADDR src, GM_ADDR y, 
                                int32_t loopCount, int32_t cycles, int32_t interval, 
                                int32_t reduce, bool include_self, int8_t formerNum, int16_t tailLength) {
        // ASSERT(GetBlockNum() != 0 && "block dim can not be zero!");
        
        this->loopCount = loopCount;
        this->cycles = cycles;
        this->interval = interval;

        this->reduce = reduce;
        this->include_self = include_self;

        if(GetBlockIdx()%2 == 0)
        {
            selfGm.SetGlobalBuffer((__gm__ DTYPE_SELF*)self, 1);
            indexGm.SetGlobalBuffer((__gm__ DTYPE_INDEX*)index, 1);
            srcGm.SetGlobalBuffer((__gm__ DTYPE_SRC*)src, 1);
            yGm.SetGlobalBuffer((__gm__ DTYPE_Y*)y, 1);
        }
        else
        {
            selfGm.SetGlobalBuffer((__gm__ DTYPE_SELF*)self+interval, 1);
            indexGm.SetGlobalBuffer((__gm__ DTYPE_INDEX*)index+interval, 1);
            srcGm.SetGlobalBuffer((__gm__ DTYPE_SRC*)src+interval, 1);
            yGm.SetGlobalBuffer((__gm__ DTYPE_Y*)y+interval, 1);
        }

        if ((GetBlockIdx()/2) < formerNum)
        {
            this->cycles_num = (tailLength + 1);
            this->cycles_start = (tailLength + 1) * (GetBlockIdx()/2);
            this->cycles_end = this->cycles_start + tailLength + 1;
        }
        else
        {
            this->cycles_num = tailLength;
            this->cycles_start = tailLength * (GetBlockIdx()/2) + formerNum;
            this->cycles_end = this->cycles_start + tailLength;
        }

        if constexpr (std::is_same_v<DTYPE_Y, half>)
        {
            this->interval_32 = (this->interval + 15)/16*16;
        }
        else
        {
            this->interval_32 = (this->interval + 7)/8*8;
        }


        

        pipe.InitBuffer(inQueueSelf, BUFFER_NUM, (this->interval)* sizeof(DTYPE_SELF));
        pipe.InitBuffer(inQueueIndex, BUFFER_NUM, (this->interval) * sizeof(DTYPE_INDEX));
        pipe.InitBuffer(inQueueSrc, BUFFER_NUM, (this->interval) * sizeof(DTYPE_SRC));
        // pipe.InitBuffer(outQueueY, 1, (this->interval_32* this->cycles_num)* sizeof(DTYPE_Y));
        pipe.InitBuffer(outQueueY, BUFFER_NUM, (this->interval) * ACCTEST_NUM * sizeof(DTYPE_Y));
        

        pipe.InitBuffer(QueueTmp1, (this->interval) * sizeof(uint8_t));
        pipe.InitBuffer(QueueTmp2, (this->interval)* sizeof(DTYPE_SRC));
    }
 
    __aicore__ inline void Process() 
    {

        int cycles_end_align = this->cycles_end / ACCTEST_NUM *ACCTEST_NUM;
        {
            auto tmp1 = QueueTmp1.Get<uint8_t>();
            for(int i=0; i<this->loopCount; i++)
            {
                for(int j=this->cycles_start; j<cycles_end_align; j+=ACCTEST_NUM)
                {
                    LocalTensor<DTYPE_Y> yLocal = outQueueY.AllocTensor<DTYPE_Y>();
                    if constexpr (std::is_same_v<DTYPE_Y, half>)
                    {
                        Duplicate(yLocal, (DTYPE_Y)(65504), this->interval*ACCTEST_NUM);
                    }
                    else
                    {
                        Duplicate(yLocal, (DTYPE_Y)(3.40282e+038), this->interval*ACCTEST_NUM);
                    }
                    
                    for(int z=0; z<this->cycles; z++)
                    {
                        LocalTensor<DTYPE_Y> srcLocal = inQueueSrc.AllocTensor<DTYPE_Y>();
                        LocalTensor<DTYPE_INDEX> indexLocal = inQueueIndex.AllocTensor<DTYPE_INDEX>();
                        DataCopy(srcLocal, srcGm[i*this->cycles*this->interval*2 + z*this->interval*2], this->interval);
                        DataCopy(indexLocal, indexGm[i*this->cycles*this->interval*2 + z*this->interval*2], this->interval);
                        inQueueSrc.EnQue(srcLocal);
                        inQueueIndex.EnQue(indexLocal);

                        LocalTensor<DTYPE_Y> srcLocal1 = inQueueSrc.DeQue<DTYPE_Y>();
                        LocalTensor<DTYPE_INDEX> indexLocal1 = inQueueIndex.DeQue<DTYPE_INDEX>();
                        auto tmp2 = QueueTmp2.Get<DTYPE_Y>();

                        for(int k=0; k<ACCTEST_NUM; k++)
                        {
                            CompareScalar(tmp1, indexLocal1, (j+k), CMPMODE::EQ, (this->interval+127)/128*128);
                            Select(tmp2, tmp1, srcLocal1, yLocal[this->interval*k], SELMODE::VSEL_TENSOR_TENSOR_MODE, this->interval);
                            Min(yLocal[this->interval*k], yLocal[this->interval*k], tmp2, this->interval);
                        }
                        inQueueSrc.FreeTensor(srcLocal1);
                        inQueueIndex.FreeTensor(indexLocal1);
                    }
                    for(int z=0; z<ACCTEST_NUM; z++)
                    {
                        LocalTensor<DTYPE_Y> selfLocal = inQueueSelf.AllocTensor<DTYPE_SELF>();
                        DataCopy(selfLocal, selfGm[i*this->cycles*this->interval*2 + (j+z)*this->interval*2], this->interval);
                        inQueueSelf.EnQue(selfLocal);
                        selfLocal = inQueueSelf.DeQue<DTYPE_SELF>();
                        if constexpr (std::is_same_v<DTYPE_Y, half>)
                        {
                            CompareScalar(tmp1, yLocal[z*this->interval], (DTYPE_Y)(65504), CMPMODE::EQ, (this->interval+127)/128*128);
                        }
                        else
                        {
                            CompareScalar(tmp1, yLocal[z*this->interval], (DTYPE_Y)(3.40282e+038), CMPMODE::EQ, (this->interval+127)/128*128);
                        }
                        Select(yLocal[z*this->interval], tmp1, selfLocal, yLocal[z*this->interval], SELMODE::VSEL_TENSOR_TENSOR_MODE, this->interval);
                        inQueueSelf.FreeTensor(selfLocal);
                    }
                    outQueueY.EnQue(yLocal);

                    yLocal = outQueueY.DeQue<DTYPE_Y>();
                    DataCopyExtParams copyParams{ACCTEST_NUM, static_cast<uint16_t>(this->interval * sizeof(DTYPE_Y)), 0, static_cast<uint16_t>(this->interval * sizeof(DTYPE_Y)), 0}; // 结构体DataCopyExtParams最后一个参数是rsv保留位
                    DataCopyPad(yGm[i*this->cycles*this->interval*2 + j*this->interval*2], yLocal, copyParams); 
                    outQueueY.FreeTensor(yLocal);
                }
                if(cycles_end_align != this->cycles_end)
                {
                    uint16_t acc_num =  this->cycles_end - cycles_end_align;
                    int j=cycles_end_align;
                    LocalTensor<DTYPE_Y> yLocal = outQueueY.AllocTensor<DTYPE_Y>();
                    if constexpr (std::is_same_v<DTYPE_Y, half>)
                    {
                        Duplicate(yLocal, (DTYPE_Y)(65504), this->interval*acc_num);
                    }
                    else
                    {
                        Duplicate(yLocal, (DTYPE_Y)(3.40282e+038), this->interval*acc_num);
                    }
                    
                    for(int z=0; z<this->cycles; z++)
                    {
                        LocalTensor<DTYPE_Y> srcLocal = inQueueSrc.AllocTensor<DTYPE_Y>();
                        LocalTensor<DTYPE_INDEX> indexLocal = inQueueIndex.AllocTensor<DTYPE_INDEX>();
                        DataCopy(srcLocal, srcGm[i*this->cycles*this->interval*2 + z*this->interval*2], this->interval);
                        DataCopy(indexLocal, indexGm[i*this->cycles*this->interval*2 + z*this->interval*2], this->interval);
                        inQueueSrc.EnQue(srcLocal);
                        inQueueIndex.EnQue(indexLocal);

                        LocalTensor<DTYPE_Y> srcLocal1 = inQueueSrc.DeQue<DTYPE_Y>();
                        LocalTensor<DTYPE_INDEX> indexLocal1 = inQueueIndex.DeQue<DTYPE_INDEX>();
                        auto tmp2 = QueueTmp2.Get<DTYPE_Y>();

                        for(int k=0; k<acc_num; k++)
                        {
                            CompareScalar(tmp1, indexLocal1, (j+k), CMPMODE::EQ, (this->interval+127)/128*128);
                            Select(tmp2, tmp1, srcLocal1, yLocal[this->interval*k], SELMODE::VSEL_TENSOR_TENSOR_MODE, this->interval);
                            Min(yLocal[this->interval*k], yLocal[this->interval*k], tmp2, this->interval);
                        }
                        inQueueSrc.FreeTensor(srcLocal1);
                        inQueueIndex.FreeTensor(indexLocal1);
                    }
                    for(int z=0; z<acc_num; z++)
                    {
                        LocalTensor<DTYPE_Y> selfLocal = inQueueSelf.AllocTensor<DTYPE_SELF>();
                        DataCopy(selfLocal, selfGm[i*this->cycles*this->interval*2 + (j+z)*this->interval*2], this->interval);
                        inQueueSelf.EnQue(selfLocal);
                        selfLocal = inQueueSelf.DeQue<DTYPE_SELF>();
                        if constexpr (std::is_same_v<DTYPE_Y, half>)
                        {
                            CompareScalar(tmp1, yLocal[z*this->interval], (DTYPE_Y)(65504), CMPMODE::EQ, (this->interval+127)/128*128);
                        }
                        else
                        {
                            CompareScalar(tmp1, yLocal[z*this->interval], (DTYPE_Y)(3.40282e+038), CMPMODE::EQ, (this->interval+127)/128*128);
                        }
                        Select(yLocal[z*this->interval], tmp1, selfLocal, yLocal[z*this->interval], SELMODE::VSEL_TENSOR_TENSOR_MODE, this->interval);
                        inQueueSelf.FreeTensor(selfLocal);
                    }
                    outQueueY.EnQue(yLocal);

                    yLocal = outQueueY.DeQue<DTYPE_Y>();
                    DataCopyExtParams copyParams{acc_num, static_cast<uint16_t>(this->interval * sizeof(DTYPE_Y)), 0, static_cast<uint16_t>(this->interval * sizeof(DTYPE_Y)), 0}; // 结构体DataCopyExtParams最后一个参数是rsv保留位
                    DataCopyPad(yGm[i*this->cycles*this->interval*2 + j*this->interval*2], yLocal, copyParams); 
                    outQueueY.FreeTensor(yLocal);
                }
            }
        }
    }

private:
    __aicore__ inline void CopyIn1(int32_t progress) {

        DataCopyExtParams copyParams{1, static_cast<uint16_t>(this->interval * sizeof(DTYPE_SELF)), 0, 0, 0}; // 结构体DataCopyExtParams最后一个参数是rsv保留位
        DataCopyPadExtParams<DTYPE_Y> padParams{false, 0, 0, 0};
        LocalTensor<DTYPE_Y> selfLocal = QueueSelf.AllocTensor<DTYPE_SELF>();
        DataCopyPad(selfLocal, selfGm[progress * this->interval], copyParams, padParams); 
        QueueSelf.EnQue(selfLocal);
    }
    __aicore__ inline void CopyOut1(int32_t progress) {
        DataCopyExtParams copyParams{1, static_cast<uint16_t>(this->interval * sizeof(DTYPE_SELF)), 0, 0, 0}; // 结构体DataCopyExtParams最后一个参数是rsv保留位
        LocalTensor<DTYPE_SELF> selfLocal = QueueSelf.DeQue<DTYPE_SELF>();
        DataCopyPad(yGm[progress * this->interval], selfLocal, copyParams); 
        QueueSelf.FreeTensor(selfLocal);
    }

private:
    TPipe pipe;

    GlobalTensor<DTYPE_SELF> selfGm;
    GlobalTensor<DTYPE_INDEX> indexGm;
    GlobalTensor<DTYPE_SRC> srcGm;
    GlobalTensor<DTYPE_Y> yGm;

    TQueBind<QuePosition::VECIN, QuePosition::VECOUT, BUFFER_NUM> QueueSelf;
    TQue<QuePosition::VECIN, BUFFER_NUM> inQueueSelf, inQueueIndex, inQueueSrc;
    TQue<QuePosition::VECOUT, BUFFER_NUM> outQueueY;
    TBuf<QuePosition::VECCALC> QueueTmp1, QueueTmp2,QueueTmp3,QueueTmp4, QueueTmp5;

    int32_t loopCount;
    int32_t cycles;
    int32_t interval;

    int32_t interval_32;

    int32_t cycles_start;
    int32_t cycles_num;
    int32_t cycles_end;
    
    int32_t reduce;
    bool include_self;
};

class KernelScatterReduce_all {
public:
    __aicore__ inline KernelScatterReduce_all() {}
    __aicore__ inline void Init(GM_ADDR self, GM_ADDR index, GM_ADDR src, GM_ADDR y, 
                                int32_t self_size, int32_t index_size, int32_t *self_ndarray, int32_t *index_ndarray, int32_t *src_ndarray,
                                int32_t index_dim, int32_t dim, int32_t reduce, bool include_self) {
        ASSERT(GetBlockNum() != 0 && "block dim can not be zero!");
        
        this->self_size = self_size;
        this->index_size = index_size;
        this->index_dim = index_dim;
        for(int i=0; i<this->index_dim; i++)
        {
            this->self_ndarray[i] = self_ndarray[i];
            this->index_ndarray[i] = index_ndarray[i];
            this->src_ndarray[i] = src_ndarray[i];
        }
        this->dim = dim;
        this->reduce = reduce;
        this->include_self = include_self;
        selfGm.SetGlobalBuffer((__gm__ DTYPE_SELF*)self, 1);
        indexGm.SetGlobalBuffer((__gm__ DTYPE_INDEX*)index, 1);
        srcGm.SetGlobalBuffer((__gm__ DTYPE_SRC*)src, 1);
        yGm.SetGlobalBuffer((__gm__ DTYPE_Y*)y, 1);
    }
    __aicore__ inline void Process() {

        for (int32_t i = 0; i < this->self_size; i++)
        {
            yGm.SetValue(i, (DTYPE_Y)(selfGm.GetValue(i)));
        }
        if(this->reduce == 0) //sum
        {
            if(this->include_self == false)
            {
                for (int32_t i = 0; i < this->index_size; i++)
                {
                    int32_t result[10]; 
                    int32_t self_idx = 0;
                    int32_t src_idx = 0;

                    unravel_index(result, i);
                    result[this->dim] = indexGm.GetValue(i);

                    for (int j = 0; j < this->index_dim; ++j)
                    {
                        self_idx *= this->self_ndarray[j];
                        self_idx += result[j];
                    }

                    yGm.SetValue(self_idx, (DTYPE_Y)(0));
                }
            }
            for (int32_t i = 0; i < this->index_size; i++)
            {
                int32_t result[10]; 
                int32_t self_idx = 0;
                int32_t src_idx = 0;

                unravel_index(result, i);
                for (int j = 0; j < this->index_dim; ++j)
                {
                    src_idx *= this->src_ndarray[j];
                    src_idx += result[j];
                }

                result[this->dim] = indexGm.GetValue(i);

                for (int j = 0; j < this->index_dim; ++j)
                {
                    self_idx *= this->self_ndarray[j];
                    self_idx += result[j];
                }

                float y_var = (float)yGm.GetValue(self_idx);
                float src_var = (float)srcGm.GetValue(src_idx);
                yGm.SetValue(self_idx, (DTYPE_Y)(y_var + src_var));
            }
        }
        else if(this->reduce == 1) //prod
        {
            if(this->include_self == false)
            {
                for (int32_t i = 0; i < this->index_size; i++)
                {
                    int32_t result[10]; 
                    int32_t self_idx = 0;
                    int32_t src_idx = 0;

                    unravel_index(result, i);
                    result[this->dim] = indexGm.GetValue(i);

                    for (int j = 0; j < this->index_dim; ++j)
                    {
                        self_idx *= this->self_ndarray[j];
                        self_idx += result[j];
                    }

                    yGm.SetValue(self_idx, (DTYPE_Y)(1));
                }
            }
            for (int32_t i = 0; i < this->index_size; i++)
            {
                int32_t result[10]; 
                int32_t self_idx = 0;
                int32_t src_idx = 0;

                unravel_index(result, i);
                for (int j = 0; j < this->index_dim; ++j)
                {
                    src_idx *= this->src_ndarray[j];
                    src_idx += result[j];
                }

                result[this->dim] = indexGm.GetValue(i);

                for (int j = 0; j < this->index_dim; ++j)
                {
                    self_idx *= this->self_ndarray[j];
                    self_idx += result[j];
                }

                float y_var = (float)yGm.GetValue(self_idx);
                float src_var = (float)srcGm.GetValue(src_idx);
                yGm.SetValue(self_idx, (DTYPE_Y)(y_var * src_var));
            }
        }
        else if(this->reduce == 2) //mean
        {
            if(this->include_self == false)
            {
                for (int32_t i = 0; i < this->index_size; i++)
                {
                    int32_t result[10]; 
                    int32_t self_idx = 0;
                    int32_t src_idx = 0;

                    unravel_index(result, i);
                    result[this->dim] = indexGm.GetValue(i);

                    for (int j = 0; j < this->index_dim; ++j)
                    {
                        self_idx *= this->self_ndarray[j];
                        self_idx += result[j];
                    }

                    yGm.SetValue(self_idx, (DTYPE_Y)(0));
                }
            }
            for (int32_t i = 0; i < this->index_size; i++)
            {
                int32_t result[10]; 
                int32_t self_idx = 0;
                int32_t src_idx = 0;

                unravel_index(result, i);
                for (int j = 0; j < this->index_dim; ++j)
                {
                    src_idx *= this->src_ndarray[j];
                    src_idx += result[j];
                }

                result[this->dim] = indexGm.GetValue(i);

                for (int j = 0; j < this->index_dim; ++j)
                {
                    self_idx *= this->self_ndarray[j];
                    self_idx += result[j];
                }

                float y_var = (float)yGm.GetValue(self_idx);
                float src_var = (float)srcGm.GetValue(src_idx);
                yGm.SetValue(self_idx, (DTYPE_Y)(y_var * src_var));
            }
        }
        else if(this->reduce == 3) //amax
        {
            if(this->include_self == false)
            {
                for (int32_t i = 0; i < this->index_size; i++)
                {
                    int32_t result[10]; 
                    int32_t self_idx = 0;
                    int32_t src_idx = 0;

                    unravel_index(result, i);
                    result[this->dim] = indexGm.GetValue(i);

                    for (int j = 0; j < this->index_dim; ++j)
                    {
                        self_idx *= this->self_ndarray[j];
                        self_idx += result[j];
                    }

                    if constexpr (std::is_same_v<DTYPE_Y, half>)
                    {
                        yGm.SetValue(self_idx, (DTYPE_Y)(-65504));
                    }
                    else
                    {
                        yGm.SetValue(self_idx, (DTYPE_Y)(-3.40282e+038));
                    }
                }
            }
            for (int32_t i = 0; i < this->index_size; i++)
            {
                int32_t result[10]; 
                int32_t self_idx = 0;
                int32_t src_idx = 0;

                unravel_index(result, i);
                for (int j = 0; j < this->index_dim; ++j)
                {
                    src_idx *= this->src_ndarray[j];
                    src_idx += result[j];
                }

                result[this->dim] = indexGm.GetValue(i);

                for (int j = 0; j < this->index_dim; ++j)
                {
                    self_idx *= this->self_ndarray[j];
                    self_idx += result[j];
                }

                float y_var = (float)yGm.GetValue(self_idx);
                float src_var = (float)srcGm.GetValue(src_idx);
                if(y_var < src_var)
                {
                    yGm.SetValue(self_idx, (DTYPE_Y)(src_var));
                }
            }
        }
        else //amin
        {
            if(this->include_self == false)
            {
                for (int32_t i = 0; i < this->index_size; i++)
                {
                    int32_t result[10]; 
                    int32_t self_idx = 0;
                    int32_t src_idx = 0;

                    unravel_index(result, i);
                    result[this->dim] = indexGm.GetValue(i);

                    for (int j = 0; j < this->index_dim; ++j)
                    {
                        self_idx *= this->self_ndarray[j];
                        self_idx += result[j];
                    }
                    if constexpr (std::is_same_v<DTYPE_Y, half>)
                    {
                        yGm.SetValue(self_idx, (DTYPE_Y)(65504));
                    }
                    else
                    {
                        yGm.SetValue(self_idx, (DTYPE_Y)(3.40282e+038));
                    }
                }
            }
            for (int32_t i = 0; i < this->index_size; i++)
            {
                int32_t result[10]; 
                int32_t self_idx = 0;
                int32_t src_idx = 0;

                unravel_index(result, i);
                for (int j = 0; j < this->index_dim; ++j)
                {
                    src_idx *= this->src_ndarray[j];
                    src_idx += result[j];
                }

                result[this->dim] = indexGm.GetValue(i);

                for (int j = 0; j < this->index_dim; ++j)
                {
                    self_idx *= this->self_ndarray[j];
                    self_idx += result[j];
                }

                float y_var = (float)yGm.GetValue(self_idx);
                float src_var = (float)srcGm.GetValue(src_idx);
                if(y_var > src_var)
                {
                    yGm.SetValue(self_idx, (DTYPE_Y)(src_var));
                }
            }
        }
    }
private:
    __aicore__ inline void unravel_index(int32_t *result, int32_t ind)
    {
        for (int i = this->index_dim - 1; i >= 0; --i)
        {
            result[i] = ind % this->index_ndarray[i];
            ind /= index_ndarray[i];
        }
    }
private:

    TPipe pipe;

    GlobalTensor<DTYPE_SELF> selfGm;
    GlobalTensor<DTYPE_INDEX> indexGm;
    GlobalTensor<DTYPE_SRC> srcGm;
    GlobalTensor<DTYPE_Y> yGm;


    int32_t self_size;
    int32_t index_size;
    int32_t self_ndarray[8];
    int32_t index_ndarray[8];
    int32_t src_ndarray[8];

    int32_t index_dim;
    int32_t dim;
    int32_t reduce;
    bool include_self;

   
};
extern "C" __global__ __aicore__ void scatter_reduce(GM_ADDR self, GM_ADDR index, GM_ADDR src, GM_ADDR y, GM_ADDR workspace, GM_ADDR tiling) {
    GET_TILING_DATA(tiling_data, tiling);
    // TODO: user kernel impl
    if (TILING_KEY_IS(1)) {
        KernelScatterReduce op;
        op.Init(self, index, src, y,
            tiling_data.loopCount, tiling_data.cycles, tiling_data.interval,
            tiling_data.reduce, tiling_data.include_self, tiling_data.formerNum, tiling_data.tailLength);  
        op.Process();
    }
    else if (TILING_KEY_IS(2)) {
        KernelScatterReduce_all op;
        op.Init(self, index, src, y,
            tiling_data.self_size, tiling_data.index_size, tiling_data.self_ndarray, tiling_data.index_ndarray, tiling_data.src_ndarray,  
            tiling_data.index_dim, tiling_data.dim, tiling_data.reduce, tiling_data.include_self);  
        op.Process();
    }
}